import os.path

import pandas as pd


class DataOperation:
    def __init__(self, file_path):
        # 分割文件路径名，判断使用不同的方法读取文件
        _, extension = os.path.split(file_path)
        if extension == '.txt' or '.csv' or '.tsv':
            self.data = pd.read_csv(file_path)
        elif extension == '.xls' or '.xlsx':
            self.data = pd.read_excel(file_path, sheet_name=0)
        else:
            print('请输入正确的路径文件名')

    def get_head(self, num=5):
        # 读取n行数据
        return self.data.head(num)

    def get_shape(self):
        # 获取数据的行，列
        return self.data.shape

    def get_columns(self):
        # 获取数据的字段名
        return self.data.columns

    def get_index(self):
        # 获取数据的索引
        return self.data.index

    def get_data_to_dict(self):
        # 把数据转化为字典
        return [self.data.loc[i].to_dict() for i in self.data.index.values]

    def get_data_to_list(self):
        # 把数据转化为列表
        return self.data.values.tolist()

if __name__ == '__main__':
    pass


